import torch
from torch import nn
from mlp_module import MLP

class CraftModule_old(nn.Module):
    def __init__(self, input_size=29, hidden_size=58, output_size=29):
        super(CraftModule_old, self).__init__()
        self.mlp = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, output_size)
        )
        
    def forward(self, x):
        x = self.mlp(x)
        return x
    
class CraftModule(nn.Module):
    def __init__(self, input_size=29, hidden_size=58, output_size=29):
        super(CraftModule, self).__init__()
        self.mlp = MLP(input_size, hidden_size, output_size)
        
    def forward(self, x):
        x = self.mlp(x)
        return x
    
if __name__ == "__main__":
    b_s = 16
    n_feat = 30
    x = torch.randn(b_s, n_feat)
    model = CraftModule(n_feat, n_feat*2, 1)
    output = model(x)
    print(output.shape)